"""
SUMO simulation specific helper methods
"""

__author__ = "Peter Kocsis, Edmond Irani Liu"
__copyright__ = "TUM Cyber-Physical System Group"
__credits__ = []
__version__ = "0.1"
__maintainer__ = "Edmond Irani Liu"
__email__ = "edmond.irani@tum.de"
__status__ = "Integration"

import copy
import os
import pickle
from enum import unique, Enum
from math import sin, cos
from typing import Tuple, Dict, Optional

import numpy as np
from commonroad.common.file_reader import CommonRoadFileReader
from commonroad.common.solution import Solution
from commonroad.planning.planning_problem import PlanningProblemSet
from commonroad.scenario.scenario import Scenario

# sumocr 是 commonroad 与 sumo 的封装
from sumocr.interface.ego_vehicle import EgoVehicle
from sumocr.interface.sumo_simulation import SumoSimulation
from sumocr.maps.sumo_scenario import ScenarioWrapper
from sumocr.sumo_config.default import DefaultConfig
from sumocr.sumo_docker.interface.docker_interface import SumoInterface
from sumocr.visualization.video import create_video


@unique
class SimulationOption(Enum):
    WITHOUT_EGO = "_without_ego"
    MOTION_PLANNER = "_planner"
    SOLUTION = "_solution"


def simulate_scenario(mode: SimulationOption,
                      conf: DefaultConfig,
                      scenario_wrapper: ScenarioWrapper,
                      scenario_path: str,
                      num_of_steps: int = None,
                      planning_problem_set: PlanningProblemSet = None,
                      solution: Solution = None,
                      use_sumo_manager: bool = False) -> Tuple[Scenario, Dict[int, EgoVehicle]]:
    """
    Simulates an interactive scenario with specified mode

    :param mode: 仿真的类型 0 = without ego, 1 = with plugged in planner, 2 = with solution trajectory
    :param conf: 仿真使用的配置对象 config of the simulation
    :param scenario_wrapper: sumocr 中的 ScenarioWrapper 对象 scenario wrapper used by the Simulator
    :param scenario_path: commonroad 中场景文件路径 path to the interactive scenario folder
    :param num_of_steps: 仿真的步数 number of steps to simulate
    :param planning_problem_set: 使用 cr 的 filereader 读到的 planning_problem_set 对象 planning problem set of the scenario
    :param solution: 规划问题的解决对象 solution to the planning problem
    :param use_sumo_manager: 是否使用 sumo manager indicates whether to use the SUMO Manager
    :return: simulated scenario and dictionary with items {planning_problem_id: EgoVehicle}
    """
    
    # 这一步
    res = scenario_wrapper.get_rou_file() # applies corrections to route file if necessary # TODO 临时注释
    print('scenario_wrapper.get_rou_file() = {0}'.format(res))

    # 设置步长
    if num_of_steps is None:
        num_of_steps = conf.simulation_steps
        print('仿真步长是：{0}'.format(num_of_steps))

    # 设置是否使用 sumo interface
    sumo_interface = None
    if use_sumo_manager:
        sumo_interface = SumoInterface(use_docker=True)
        sumo_sim = sumo_interface.start_simulator()

        sumo_sim.send_sumo_scenario(conf.scenario_name,
                                    scenario_path)
    else:
        sumo_sim = SumoSimulation()

    # initialize simulation
    
    sumo_sim.initialize(conf, scenario_wrapper, None)

    if mode is SimulationOption.WITHOUT_EGO:
        # simulation without ego vehicle
        for step in range(num_of_steps):
            # set to dummy simulation
            sumo_sim.dummy_ego_simulation = True
            sumo_sim.simulate_step()

    elif mode is SimulationOption.MOTION_PLANNER:
        # simulation with plugged in planner
        # 如果使用带运动规划的仿真
        
        # 定义一个 run_simulation 的函数
        def run_simulation():
            ego_vehicles = sumo_sim.ego_vehicles
            for step in range(num_of_steps):
                if use_sumo_manager:
                    ego_vehicles = sumo_sim.ego_vehicles

                # retrieve the CommonRoad scenario at the current time step, e.g. as an input for a prediction module
                current_scenario = sumo_sim.commonroad_scenario_at_time_step(sumo_sim.current_time_step)    # 获得本时间戳的 快照
                for idx, ego_vehicle in enumerate(ego_vehicles.values()):
                    # retrieve the current state of the ego vehicle
                    state_current_ego = ego_vehicle.current_state    # 这里获得当前这个 step 中 current ego 的状态 类型是 CustomStep

                    # ====== plug in your motion planner here  这里填写规划算法
                    # example motion planner which decelerates to full stop 这里是一个简单的停车算法
                    next_state = copy.deepcopy(state_current_ego)
                    next_state.steering_angle = 0.0   # 下一时刻的方向盘转角为 0
                    a = -1.0    # 减速度为 -4
                    dt = 0.1    # 仿真 dt 为 0.1
                    if next_state.velocity > 0:
                        v = next_state.velocity
                        x, y = next_state.position
                        o = next_state.orientation

                        next_state.position = np.array([x + v * cos(o) * dt, y + v * sin(o) * dt])
                        next_state.velocity += a * dt
                    # ====== end of motion planner

                    # update the ego vehicle with new trajectory with only 1 state for the current step
                    next_state.time_step = 1     # 下一个时间戳为 1
                    trajectory_ego = [next_state]    # 
                    ego_vehicle.set_planned_trajectory(trajectory_ego)

                if use_sumo_manager:
                    # set the modified ego vehicles to synchronize in case of using sumo_docker
                    sumo_sim.ego_vehicles = ego_vehicles

                sumo_sim.simulate_step()

        run_simulation()

    elif mode is SimulationOption.SOLUTION:
        # simulation with given solution trajectory

        def run_simulation():
            ego_vehicles = sumo_sim.ego_vehicles

            for time_step in range(num_of_steps):
                if use_sumo_manager:
                    ego_vehicles = sumo_sim.ego_vehicles
                for idx_ego, ego_vehicle in enumerate(ego_vehicles.values()):
                    # update the ego vehicles with solution trajectories
                    trajectory_solution = solution.planning_problem_solutions[idx_ego].trajectory
                    next_state = copy.deepcopy(trajectory_solution.state_list[time_step])

                    next_state.time_step = 1
                    trajectory_ego = [next_state]
                    ego_vehicle.set_planned_trajectory(trajectory_ego)

                if use_sumo_manager:
                    # set the modified ego vehicles to synchronize in case of using SUMO Manager
                    sumo_sim.ego_vehicles = ego_vehicles

                sumo_sim.simulate_step()

        check_trajectories(solution, planning_problem_set, conf)
        run_simulation()

    # retrieve the simulated scenario in CR format
    simulated_scenario = sumo_sim.commonroad_scenarios_all_time_steps()

    # stop the simulation
    sumo_sim.stop()
    ego_vehicles = {list(planning_problem_set.planning_problem_dict.keys())[0]:
                        ego_v for _, ego_v in sumo_sim.ego_vehicles.items()}

    if use_sumo_manager:
        sumo_interface.stop_simulator()

    return simulated_scenario, ego_vehicles


def simulate_without_ego(interactive_scenario_path: str,
                         output_folder_path: str = None,
                         create_video: bool = False,
                         use_sumo_manager: bool = False,
                         num_of_steps=None) -> Tuple[Scenario, PlanningProblemSet]:
    """
    Simulates an interactive scenario without ego vehicle

    :param interactive_scenario_path: path to the interactive scenario folder 交互场景的文件夹
    :param output_folder_path: path to the output folder   仿真结果输出的文件夹
    :param create_video: indicates whether to create a mp4 of the simulated scenario    是否创建 视频
    :param use_sumo_manager: indicates whether to use the SUMO Manager   是否使用 sumo manager
    :param num_of_steps: max. number of simulated time steps    仿真多少步
    :return: Tuple of the simulated scenario and the planning problem set     返回仿真的场景和 planning problem
    """
    
    # 获得 sumo 的配置文件
    print('interactive_scenario_path = {0}'.format(interactive_scenario_path))
    conf = load_sumo_configuration(interactive_scenario_path)
    
    # 读取 cr.xml 格式的场景文件，生成场景和规划问题
    scenario_file = os.path.join(interactive_scenario_path, f"{conf.scenario_name}.cr.xml")
    print(scenario_file)
    scenario, planning_problem_set = CommonRoadFileReader(scenario_file).open()
    
    # ScenarioWrapper 是 sumocr 中的 场景封装器对象
    scenario_wrapper = ScenarioWrapper()
    # sumo.cfg 文件是 sumo 仿真器中使用的 sumo xml 文件
    scenario_wrapper.sumo_cfg_file = os.path.join(interactive_scenario_path, f"{conf.scenario_name}.sumo.cfg")
    # ScenarioWapper 除了要添加 sumo 的仿真 cfg  文件，还需要设置 commonroad 的场景对象
    scenario_wrapper.initial_scenario = scenario

    # 从 Sumocr 的 config 对象中获取要仿真的步长
    num_of_steps = conf.simulation_steps if num_of_steps is None else num_of_steps
    # simulation without ego vehicle
    simulated_scenario_without_ego, _ = simulate_scenario(SimulationOption.WITHOUT_EGO, conf,
                                                          scenario_wrapper,
                                                          interactive_scenario_path,
                                                          num_of_steps=num_of_steps,
                                                          planning_problem_set=planning_problem_set,
                                                          solution=None,
                                                          use_sumo_manager=use_sumo_manager)
    print('仿真场景中 ego_vehicle = {0}'.format(_))
    simulated_scenario_without_ego.scenario_id = scenario.scenario_id

    if create_video:
        print('开始 create_video')
        create_video_for_simulation(simulated_scenario_without_ego, output_folder_path, planning_problem_set,
                                    {}, SimulationOption.WITHOUT_EGO.value)

    return simulated_scenario_without_ego, planning_problem_set


def simulate_with_solution(interactive_scenario_path: str,
                           output_folder_path: str = None,
                           solution: Solution = None,
                           create_video: bool = False,
                           use_sumo_manager: bool = False,
                           create_ego_obstacle: bool = False) \
        -> Tuple[Scenario, PlanningProblemSet, Dict[int, EgoVehicle]]:
    """
    Simulates an interactive scenario with a given solution

    :param interactive_scenario_path: 场景文件的路径
    :param output_folder_path: 结果的输出文件夹
    :param solution: 使用带着 planner 的算法运行的 结果文件，里面是 ego_car 的状态列表
    :param create_video: indicates whether to create a mp4 of the simulated scenario
    :param use_sumo_manager: indicates whether to use the SUMO Manager
    :param create_ego_obstacle: indicates whether to create obstacles as the ego vehicles
    :return: Tuple of the simulated scenario and the planning problem set
    """
    if not isinstance(solution, Solution):
        raise Exception("Solution to the planning problem is not given.")

    # 加载 配置
    conf = load_sumo_configuration(interactive_scenario_path)
    scenario_file = os.path.join(interactive_scenario_path, f"{conf.scenario_name}.cr.xml")
    scenario, planning_problem_set = CommonRoadFileReader(scenario_file).open()

    # 场景对象封装
    scenario_wrapper = ScenarioWrapper()
    sumo_cfg_file = os.path.join(interactive_scenario_path, f"{conf.scenario_name}.sumo.cfg")
    scenario_wrapper.initialize(conf.scenario_name, sumo_cfg_file, scenario_file)
    scenario_with_solution, ego_vehicles = simulate_scenario(SimulationOption.SOLUTION, conf,
                                                             scenario_wrapper,
                                                             interactive_scenario_path,
                                                             num_of_steps=conf.simulation_steps,
                                                             planning_problem_set=planning_problem_set,
                                                             solution=solution,
                                                             use_sumo_manager=use_sumo_manager)
    scenario_with_solution.scenario_id = scenario.scenario_id

    if create_video:
        create_video_for_simulation(scenario_with_solution, output_folder_path, planning_problem_set,
                                    ego_vehicles, SimulationOption.SOLUTION.value)

    if create_ego_obstacle:
        for pp_id, planning_problem in planning_problem_set.planning_problem_dict.items():
            obstacle_ego = ego_vehicles[pp_id].get_dynamic_obstacle()
            scenario_with_solution.add_objects(obstacle_ego)

    return scenario_with_solution, planning_problem_set, ego_vehicles


def simulate_with_planner(interactive_scenario_path: str,
                          output_folder_path: str = None,
                          create_video: bool = False,
                          use_sumo_manager: bool = False,
                          create_ego_obstacle: bool = False) \
        -> Tuple[Scenario, PlanningProblemSet, Dict[int, EgoVehicle]]:
    """
    Simulates an interactive scenario with a plugged in motion planner

    :param interactive_scenario_path: 在 commonroad 官网上下载的场景文件解压后的路径
    :param output_folder_path: 仿真后的结果文件
    :param create_video: 是否创建视频
    :param use_sumo_manager: 是否使用 sumo manager，如果设置 True, 要登录 lrz 的 dockerhub
    :param create_ego_obstacle: indicates whether to create obstacles from the planned trajectories as the ego vehicles
    :return: Tuple of the simulated scenario, planning problem set, and list of ego vehicles
    """
    # 1. conf DefaultConfig 类描述的 sumo 仿真配置
    conf = load_sumo_configuration(interactive_scenario_path)
    # scenario_file cr 格式的场景文件
    scenario_file = os.path.join(interactive_scenario_path, f"{conf.scenario_name}.cr.xml")
    # scenario 描述了路网，障碍物，交通标志
    # planning_problem_set 描述了 ego car 的初始位置和目标，目标是一系列的
    scenario, planning_problem_set = CommonRoadFileReader(scenario_file).open()
    # 新建一个 ScenarioWrapper() 的对象
    scenario_wrapper = ScenarioWrapper()
    scenario_wrapper.sumo_cfg_file = os.path.join(interactive_scenario_path, f"{conf.scenario_name}.sumo.cfg")
    scenario_wrapper.initial_scenario = scenario

    scenario_with_planner, ego_vehicles = simulate_scenario(SimulationOption.MOTION_PLANNER, conf,
                                                            scenario_wrapper,
                                                            interactive_scenario_path,
                                                            num_of_steps=conf.simulation_steps,
                                                            planning_problem_set=planning_problem_set,
                                                            use_sumo_manager=use_sumo_manager)
    scenario_with_planner.scenario_id = scenario.scenario_id

    if create_video:
        create_video_for_simulation(scenario_with_planner, output_folder_path, planning_problem_set,
                                    ego_vehicles, SimulationOption.MOTION_PLANNER.value)

    if create_ego_obstacle:
        for pp_id, planning_problem in planning_problem_set.planning_problem_dict.items():
            obstacle_ego = ego_vehicles[pp_id].get_dynamic_obstacle()
            scenario_with_planner.add_objects(obstacle_ego)

    return scenario_with_planner, planning_problem_set, ego_vehicles


def load_sumo_configuration(interactive_scenario_path: str) -> DefaultConfig:
    with open(os.path.join(interactive_scenario_path, "simulation_config.p"), "rb") as input_file:
        conf = pickle.load(input_file)
    
    assert isinstance(conf, DefaultConfig)
    print('load_sumo_configuration return conf type is {0}'.format(type(conf)))
    # 这里的 conf 之所以类型是 sumocr.sumo_config.default.DefaultConfig 是因为之前已经将一个对象 保存过了
    return conf


def check_trajectories(solution: Solution, pps: PlanningProblemSet, config: DefaultConfig):
    assert len(set(solution.planning_problem_ids) - set(pps.planning_problem_dict.keys())) == 0, \
        f"Provided solution trajectories with IDs {solution.planning_problem_ids} don't match " \
        f"planning problem IDs{list(pps.planning_problem_dict.keys())}"

    for s in solution.planning_problem_solutions:
        if s.trajectory.final_state.time_step < config.simulation_steps:
            raise ValueError(f"The simulation requires {config.simulation_steps} "
                             f"states, but the solution only provides"
                             f"{s.trajectory.final_state.time_step} time steps!")


def create_video_for_simulation(scenario_with_planner: Scenario, output_folder_path: str,
                                planning_problem_set: PlanningProblemSet,
                                ego_vehicles: Optional[Dict[int, EgoVehicle]],
                                suffix: str, follow_ego: bool = True):
    """Creates the mp4 animation for the simulation result."""
    if not output_folder_path:
        print("Output folder not specified, skipping mp4 generation.")
        return

    # create mp4 animation
    print('create mp4 animation, output_folder_path = {0}'.format(output_folder_path))
    create_video(scenario_with_planner,
                 output_folder_path,
                 planning_problem_set=planning_problem_set,
                 trajectory_pred=ego_vehicles,
                 follow_ego=follow_ego,
                 suffix=suffix)
